#include <algorithm>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <iterator>
#include <list>
#include <map>
#include <set>
#include <sstream>
#include <string>
#include <queue>
#include <vector>
#include <cmath>
#include <boost/program_options.hpp>
#include "MetastableSet.h"
#include "StringSplit.h"
#include "BlockAverage.h"
#include "Interpolation.h"
#include "TrajLoader.h"
#include "Bond.h"

#define MaxLineLength 2048

namespace po = boost::program_options;
using namespace std;

inline int bondIdx (const double & myangle,
		     const double & x0,
		     const double & bin)
{
  return int((myangle - x0) / bin);
}

int main(int argc, char * argv[])
{
  std::string ifile, ofile, jfile, tfile;
  double center, width, bin, sigma, epsilon, bondDx, recycleDt, x0, x1;
  unsigned nBlock;
  double beta = 1./2.4923;

  po::options_description desc ("Allow options");
  desc.add_options()
    ("help,h", "print this message")
    ("meta-center,c",  po::value<double > (&center)->default_value (180),   "center of the metastable set")
    ("meta-width,w",  po::value<double > (&width)->default_value (15),   "width of the metastable set")
    ("epsilon,e",  po::value<double > (&epsilon)->default_value (1.),   "epsilon")
    ("sigma,s",  po::value<double > (&sigma)->default_value (1.),   "sigma")
    ("bond-dx,d",  po::value<double > (&bondDx)->default_value (0.001),   "the dx to cal grad bond")
    ("recycle-dt",  po::value<double > (&recycleDt)->default_value (0.01),   "time interval of reuse the traj")
    ("input-j,j", po::value<std::string > (&jfile)->default_value ("table.j.xvg"), "table of J")
    ("x0,l",  po::value<double > (&x0)->default_value (0.0),   "the start of the range calculate J")
    ("x1,u",  po::value<double > (&x1)->default_value (1.5),   "the end of the range calculate J")
    ("bin,r",  po::value<double > (&bin)->default_value (5.),   "the bin size")
    ("num-block,n",  po::value<unsigned > (&nBlock)->default_value (10),   "number of block")
    ("input-bond,f", po::value<std::string > (&ifile)->default_value ("angaver.name"), "file including bond file names")
    ("input-traj,t", po::value<std::string > (&tfile)->default_value ("traj.name"), "file including traj file names")
    ("output,o", po::value<std::string > (&ofile)->default_value ("fmpt.out"), "the output of first mean passage time");
  
      
  po::variables_map vm;
  po::store(po::parse_command_line(argc, argv, desc), vm);
  po::notify (vm);
  if (vm.count("help")){
    std::cout << desc<< "\n";
    return 0;
  }

  unsigned nbin = (x1 - x0 + 0.01 * bin) / bin;
  bin = (x1 - x0) / double(nbin);
  cout << "# bin size: " << bin << endl;
  cout << "# nbin: " << nbin << endl;
  // record
  vector<vector<double > > record (nbin);
  vector<vector<double > > record_j (nbin);
  vector<vector<double > > record_lap (nbin);
  vector<vector<double > > record_exp (nbin);
  vector<double > sum_exp (nbin, 0.);
  vector<int > count_exp (nbin, 0);
  cout << "# j table is " << jfile << endl;
  cout << "# bond names is " << ifile << endl;
  cout << "# traj names  is " << tfile << endl;

  double metal, metau;
  metal = center - width;
  metau = center + width;
  cout << "# metastable set: [ " << metal << " , " << metau << " ]" << endl;
  // metastable set
  MetastableSet ms (metal, metau);

  // table of j
  vector<double > xx0, vv0, dd0;
  {
    ifstream fpname0 (jfile.c_str());
    if (!fpname0){
      cerr << "cannot open file " << jfile << endl;
      return 1;
    }
    char nameline [MaxLineLength];
    while (fpname0.getline(nameline, MaxLineLength)){
      if (nameline[0] == '#' || nameline[0] == '@') continue;
      vector<string > words;
      StringOperation::split (string(nameline), words);
      if (words.size() < 3) {
	cerr << "wrong file format of " << jfile << endl;;
	exit (1);
      }
      xx0.push_back (atof(words[0].c_str()));
      vv0.push_back (atof(words[1].c_str()));
      dd0.push_back (atof(words[2].c_str()));
    }
  }
  PiecewisePoly pp;
  Interpolation::piecewiseLinear (xx0, dd0, pp);
  
  // analyze traj
  ifstream fpname (ifile.c_str());
  ifstream fpname_t (tfile.c_str());
  char nameline [MaxLineLength];
  char nameline_t [MaxLineLength];
  if (!fpname){
    std::cerr << "cannot open file " << ifile << std::endl;
    return 1;
  }
  if (!fpname_t){
    std::cerr << "cannot open file " << tfile << std::endl;
    return 1;
  }
  unsigned countFound = 0;
  unsigned countUnFound = 0;
  unsigned countFile = 0;
  unsigned printCount = 0;
  
  while (fpname.getline(nameline, MaxLineLength)){
    if (nameline[0] == '#') continue;
    if (!fpname_t.getline(nameline_t, MaxLineLength)){
      cerr << "inconsistent bond and traj name files... exit" << endl;
      exit(1);
    }
    if (nameline_t[0] == '#') continue;    
    ifstream angname (nameline);
    if (!angname){
      std::cerr << "cannot open file " << nameline << std::endl;
      return 1;
    }
    TrajLoader_xtc traj_l_xtc;
    traj_l_xtc.reinit (nameline_t);
    // cout << "reading file " << nameline << endl;
    if (printCount == 1) {
      printf ("# reading file %s       \r", nameline);
      fflush (stdout);
      printCount = 0;
    }
    printCount ++;
    countFile ++;
    vector<double > times;
    vector<double > bondv;
    vector<vector<vector<double > > > c2Coord;
    vector<vector<double > > boxes;
    char valueline [MaxLineLength];
    while (angname.getline(valueline, MaxLineLength)){
      if (valueline[0] == '#' || valueline[0] == '@') continue;
      vector<string > words;
      StringOperation::split (string(valueline), words);
      if (words.size() < 2) {
	cerr << "wrong file format of " << ifile << endl;
	exit (1);
      }
      times.push_back (atof(words[0].c_str()));
      bondv.push_back (atof(words[1].c_str()));
      if (!traj_l_xtc.load()){
	cerr << "inconsistent bond and traj files... exit" << endl;
	exit(1);
      }
      if (fabs(traj_l_xtc.getTime() - times.back()) > 1e-4){
	cerr << "inconsistent times in bond and traj fiels.. exit" << endl;
	exit(1);
      }
      boxes.push_back (traj_l_xtc.getBox());
      vector<vector<double > > tmpc2;
      traj_l_xtc.formCoords (tmpc2);
      c2Coord.push_back(tmpc2);
      BondCalculator ac (boxes.back());
      if (fabs(bondv.back() -  ac.calBond (c2Coord.back())) > 1e-3) {
	cerr << "bond v and calculated bond diffs more than 1e-3: "
	     << bondv.back()
	     << " and "
	     << ac.calBond (c2Coord.back())
	     << endl;
      }
    }
    
    unsigned recycleDN = unsigned (recycleDt / (times[1] - times[0]));
    double sum_j = 0.;
    double sum_lap = 0.;
    if (ms.inSet(bondv[0])) {
      record[bondIdx(bondv[0], x0, bin)].push_back (0.0);
      record_j[bondIdx(bondv[0], x0, bin)].push_back (0.0);
      record_lap[bondIdx(bondv[0], x0, bin)].push_back (0.0);
      record_exp[bondIdx(bondv[0], x0, bin)].push_back (1.0);
      countFound ++;
    }
    else {
      bool found = true;
      unsigned nextstab = 0;
      do {
	nextstab ++;
	if (nextstab >= bondv.size()){
	  found = false;
	  // cout << "#@ bond " << nameline << " does not hit" << endl;
	  break;
	}
      } while (! ms.inSet(bondv[nextstab]));
      if (found) {
	countFound ++;
	record [bondIdx(bondv[0], x0, bin)].push_back (-(times[0] - times[nextstab]) * epsilon);
	record_exp [bondIdx(bondv[0 ], x0, bin)].push_back (exp((times[0 ] - times[nextstab]) * epsilon * beta));
	sum_exp [bondIdx(bondv[0 ], x0, bin)] += exp((times[0 ] - times[nextstab]) * epsilon * beta);
	count_exp [bondIdx(bondv[0 ], x0, bin)] ++;
	sum_j = 0.;
	sum_lap = 0.;
	for (int ii = nextstab - 1; ii >= 0; --ii){
	  BondCalculator ac (boxes[ii]);
	  double tmp = - pp.value_periodic (bondv[ii]) ;
	  double tmpBond = ac.calBond (c2Coord[ii]);
	  vector<vector<double > > tmpGrad;
	  ac.calGrad (c2Coord[ii], bondDx, tmpGrad);
	  double sumGrad = 0.;
	  if (tmpGrad.size() != 2){
	    cerr << "incorrect grad size!" << endl;
	  }
	  for (unsigned kk = 0; kk < tmpGrad.size(); ++kk){
	    for (unsigned ll = 0; ll < 3; ++ll){
	      sumGrad += tmpGrad[kk][ll] * tmpGrad[kk][ll];
	    }
	  }
	  if (fabs(bondv[ii] - tmpBond) > 0.01){
	    cerr << "inconsistent bond: " << tmpBond << " and " << bondv[ii] << endl;
	  }
	  double tmplap = ac.calLaplacian (c2Coord[ii], bondDx);
	  // printf ("time %f, bond %f %f, sumGrad %e, dj is %f\n", times[ii], bondv[ii], tmpBond, sumGrad, tmp);
	  sum_j += 0.5 * sigma * sigma * tmp * tmp * sumGrad * (times[1] - times[0]);
          sum_lap += 0.125 * sigma * sigma / beta / beta * (tmplap * tmplap / sumGrad / sumGrad) * (times[1] - times[0]);
	  if (recycleDN != 0 && ii != 0 && (nextstab - ii) % recycleDN == 0){
	    record   [bondIdx(bondv[ii], x0, bin)].push_back (-(times[ii] - times[nextstab]) * epsilon);
	    record_exp [bondIdx(bondv[ii], x0, bin)].push_back (exp((times[ii] - times[nextstab]) * epsilon * beta));
	    sum_exp [bondIdx(bondv[ii], x0, bin)] += exp((times[ii] - times[nextstab]) * epsilon * beta);
	    count_exp [bondIdx(bondv[ii], x0, bin)] ++;
	    record_j [bondIdx(bondv[ii], x0, bin)].push_back (sum_j);
	    record_lap [bondIdx(bondv[ii], x0, bin)].push_back (sum_lap);
	  }
	}
	record_j [bondIdx(bondv[0], x0, bin)].push_back (sum_j);
	record_lap [bondIdx(bondv[0], x0, bin)].push_back (sum_lap);
      }
      else {
	countUnFound ++;
      }
    }
  }
  printf ("\n");
  printf ("# read %d files, %d ( %.1f %% ) hit meta, %d ( %.1f %% ) do not hit\n",
	  countFile,
	  countFound, ((double)(countFound))/((double)(countFile)) * 100.,
	  countUnFound, ((double)(countUnFound))/((double)(countFile)) * 100.);

  {
    vector<double > xx, vv, vve, vv0, vv0e, vv1, vv1e, vv2, vv2e;
    vector<int > is_ok;
    BlockAverage ba;  
    BlockAverage ba_j;  
    BlockAverage ba_lap;  
    BlockAverage ba_exp;  
    for (unsigned ii = 0; ii < record.size(); ++ii){
      double myangle = x0 + bin * (ii + 0.5);
      if (record[ii].size() == 0 || record[ii].size() < nBlock){
	xx.push_back (myangle);
	vv.push_back (0.0);
	vve.push_back (0.0);
	vv0.push_back (0.0);
	vv0e.push_back (0.0);
	vv1.push_back (0.0);
	vv1e.push_back (0.0);
	vv2.push_back (0.0);
	vv2e.push_back (0.0);
	is_ok.push_back (0);
	// fprintf (fout, "%f %e %e    0. 0.    0. 0.   o\n",
	// 	       myangle, 0., 0.
	// 	       );
      }
      else {
	if (record[ii].size() < nBlock * 5){
	  ba.processData (record[ii], record[ii].size());
	  ba_j.processData (record_j[ii], record_j[ii].size());
	  ba_lap.processData (record_lap[ii], record_lap[ii].size());
	  ba_exp.processData (record_exp[ii], record_exp[ii].size());
	}
	else {
	  ba.processData (record[ii], nBlock);
	  ba_j.processData (record_j[ii], nBlock);
	  ba_lap.processData (record_lap[ii], nBlock);
	  ba_exp.processData (record_exp[ii], nBlock);
	}
	xx.push_back (myangle);
	vv.push_back (ba.getAvg() + ba_j.getAvg());
	vve.push_back (ba.getAvgError() + ba_j.getAvgError());
	vv0.push_back (ba.getAvg());
	vv0e.push_back (ba.getAvgError());
	vv1.push_back (ba_j.getAvg());
	vv1e.push_back (ba_j.getAvgError());
	vv2.push_back (ba_exp.getAvg());
	vv2e.push_back (ba_exp.getAvgError());
	is_ok.push_back (1);
	// fprintf (fout, "%f %e %e    %e %e    %e %e   i\n",
	// 	       myangle,
	// 	       ba.getAvg() + ba_j.getAvg(), ba.getAvgError() + ba_j.getAvgError(),
	// 	       ba.getAvg(), ba.getAvgError(),
	// 	       ba_j.getAvg(), ba_j.getAvgError()
	// 	       );
      }
    }

    int ok_posi = -1;
    for (unsigned ii = 0; ii < xx.size(); ++ii){
      if (is_ok[ii]){
	ok_posi = ii;
	break;
      }
    }
    if (ok_posi != -1){
      for (int ii = 0; ii < ok_posi; ++ii){
	vv[ii] = vv[ok_posi];
	vve[ii] = vve[ok_posi];
	vv0[ii] = vv0[ok_posi];
	vv0e[ii] = vv0e[ok_posi];
	vv1[ii] = vv1[ok_posi];
	vv1e[ii] = vv1e[ok_posi];
	vv2[ii] = vv2[ok_posi];
	vv2e[ii] = vv2e[ok_posi];
      }
    }
    ok_posi = -1;
    for (unsigned ii = xx.size() - 1; ii >= 0; --ii){
      if (is_ok[ii]){
	ok_posi = ii;
	break;
      }
    }
    if (ok_posi != -1){
      for (unsigned ii = xx.size() - 1; ii > unsigned(ok_posi); --ii){
	vv[ii] = vv[ok_posi];
	vve[ii] = vve[ok_posi];
	vv0[ii] = vv0[ok_posi];
	vv0e[ii] = vv0e[ok_posi];
	vv1[ii] = vv1[ok_posi];
	vv1e[ii] = vv1e[ok_posi];
	vv2[ii] = vv2[ok_posi];
	vv2e[ii] = vv2e[ok_posi];
      }
    }
      
    FILE * fout = fopen (ofile.c_str(), "w");
    if (fout == NULL){
      cout << "cannot open file " << ofile<< endl;
      exit (1);
    }
    for (unsigned ii = 0; ii < xx.size(); ++ii){
      double tmp;
      if (count_exp[ii] > 0){
	tmp = -1./beta * log(sum_exp[ii] / double(count_exp[ii]));
      }
      else {
	tmp = 0.;
      }
      double lb, ub;
      if (vv2[ii] - 2. * vv2e[ii] < 0) {
	lb = 0;
      }
      else{
	lb = -1./beta * log(vv2[ii] - 2. * vv2e[ii]);
      }
      ub = -1./beta * log(vv2[ii] + 2. * vv2e[ii]);
      fprintf (fout, "%f   %e %e   %e %e   %e %e   %e %e %e   %e\n",
	       xx[ii],
	       vv[ii], vve[ii],
	       vv0[ii], vv0e[ii],
	       vv1[ii], vv1e[ii],
	       -1./beta * log(vv2[ii]), lb, ub, 
	       tmp
	       );
    }
    fclose (fout);  
  }  
  return 0;
}
